import numpy as np
import matplotlib.pyplot as plt


def butterfly_operation(x, N, k):
    """
    执行蝶形运算

    参数:
    x -- 输入信号
    N -- 信号长度
    k -- 蝶形运算的下标

    返回:
    输出信号
    """
    even_index = k
    odd_index = k + N // 2
    twiddle_factor = np.exp(-2j * np.pi * k / N)

    even_part = x[even_index]
    odd_part = x[odd_index] * twiddle_factor

    return even_part + odd_part, even_part - odd_part


def visualize_butterfly(N):
    """
    可视化蝶形运算过程

    参数:
    N -- 信号长度
    """
    num_stages = int(np.log2(N))
    num_butterflies = N // 2

    fig, axs = plt.subplots(num_stages, num_butterflies, figsize=(12, 6))
    fig.subplots_adjust(hspace=0.6)

    x = np.zeros(N, dtype=np.complex128)
    x[0] = 1  # 输入信号设置为1

    for stage in range(num_stages):
        for butterfly in range(num_butterflies):
            ax = axs[stage, butterfly]
            ax.cla()  # 清空当前子图

            ax.set_xlim([-1.5, 1.5])
            ax.set_ylim([-1.5, 1.5])

            ax.axis('off')  # 隐藏坐标轴

            ax.text(-0.2, 0, f'{x[butterfly]:.2f}', fontsize=12)  # 输入信号值

            y1, y2 = butterfly_operation(x, N, butterfly)

            ax.text(1.2, 0.5, f'{y1:.2f}', color='r', fontsize=12)  # 输出信号值1
            ax.text(1.2, -0.5, f'{y2:.2f}', color='r', fontsize=12)  # 输出信号值2

            ax.arrow(0, 0, 0.8, 0, head_width=0.1, head_length=0.1, fc='b', ec='b')  # 输入信号箭头
            ax.arrow(1, 0.5, 0.3, 0, head_width=0.1, head_length=0.1, fc='r', ec='r')  # 输出信号箭头1
            ax.arrow(1, -0.5, 0.3, 0, head_width=0.1, head_length=0.1, fc='r', ec='r')  # 输出信号箭头2

            ax.set_title(f'Stage {stage + 1}, Butterfly {butterfly + 1}', fontsize=12)

            # 更新输入信号值
            x[butterfly], x[butterfly + N // 2] = y1, y2

    plt.suptitle('Butterfly Operation Visualization', fontsize=16)
    plt.show()


# 示例用法
N = 8  # 信号长度
visualize_butterfly(N)
